from patagent.task.confluence_retriever import sparse_embedding_func
from patagent.task import *
import json
from urllib.parse import unquote
from patagent.utils.data_util import is_empty
from patagent.utils.http_util import http_request
from patagent.model import ApiRequest


def chitchat(state):
    """
    Chitchat

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, documents, that contains retrieved documents
    """
    print("---CHITCHAT---")
    question = state["question"]

    # RAG generation
    generation = chitchat_chain.invoke({"question": question})
    return {"documents": [], "question": question, "generation": generation}

def retrieve(state):
    """
    Retrieve documents

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, documents, that contains retrieved documents
    """
    print("---RETRIEVE---")
    question = state["question"]

    # Retrieval
    documents = confluence_retriever.invoke(question)
    print(documents)
    return {"documents": documents, "question": question}

def hybrid_retrieve(state):
    """
    Retrieve documents

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, documents, that contains retrieved documents
    """
    print("---RETRIEVE---")
    question = state["question"]

    # Retrieval
    documents = hybrid_retriever.invoke(question)
    print(documents)
    return {"documents": documents, "question": question}


def generate(state):
    """
    Generate answer

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, generation, that contains LLM generation
    """
    print("---GENERATE---")
    question = state["question"]
    documents = state["documents"]

    # RAG generation
    generation = answer_chain.invoke({"context": documents, "question": question})
    if generation.endswith('。'):
        generation = generation[:-1]
    return {"documents": documents, "question": question, "generation": generation}


def grade_documents(state):
    """
    Determines whether the retrieved documents are relevant to the question.

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): Updates documents key with only filtered relevant documents
    """

    print("---CHECK DOCUMENT RELEVANCE TO QUESTION---")
    question = state["question"]
    documents = state["documents"]

    # Score each doc
    filtered_docs = []
    for d in documents:
        score = retrieval_grader.invoke(
            {"question": question, "document": d.page_content}
        )
        grade = score.binary_score
        if grade == "yes":
            print("---GRADE: DOCUMENT RELEVANT---")
            filtered_docs.append(d)
        else:
            print("---GRADE: DOCUMENT NOT RELEVANT---")
            continue
    return {"documents": filtered_docs, "question": question}


def transform_query(state):
    """
    Transform the query to produce a better question.

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): Updates question key with a re-phrased question
    """

    print("---TRANSFORM QUERY---")
    question = state["question"]
    documents = state["documents"]

    # Re-write question
    better_question = question_rewriter.invoke({"question": question})
    return {"documents": documents, "question": better_question}


def request_data(state):
    """
    Transform the query to produce a better question.

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): Updates question key with a re-phrased question
    """

    print("---REQUEST DATA---")
    question = state['question']

    # retrieve api
    api_info = api_retriever.invoke(question)

    # extract api fields
    extraction: ApiRequest = extract_chain.invoke(api_info)
    if is_empty(extraction.url) or is_empty(extraction.method) or is_empty(extraction.parameters):
        return {"question": question, "generation": '无法获取完整的OpenApi请求参数，请重新尝试'}
    
    res = http_request(extraction).json()
    generation = '```json \n' + unquote(json.dumps(res['data']).encode('utf-8').decode('unicode-escape')) + '\n ```'
    return {"question": question, "generation": generation}



### Edges ###

def route_question(state):
    """
    Route question to CHITCHAT or RAG.

    Args:
        state (dict): The current graph state

    Returns:
        str: Next node to call
    """

    print("---ROUTE QUESTION---")
    question = state["question"]
    print(question)
    source = question_router.invoke({"question": question})
    if source.datasource == "chitchat":
        print("---ROUTE QUESTION TO CHITCHAT---")
        return "chitchat"
    elif source.datasource == "data":
        print("---ROUTE QUESTION TO DATA---")
        return "data"
    elif source.datasource == "vectorstore":
        sparse_embedding = sparse_embedding_func.embed_query(question)
        if sparse_embedding == {}:
            print("---ROUTE QUESTION TO RETRIEVE---")
            return "retrieve"
        else:
            print("---ROUTE QUESTION TO HYBRID RETRIEVE---")
            return "hybrid_retrieve"


def decide_to_generate(state):
    """
    Determines whether to generate an answer, or re-generate a question.

    Args:
        state (dict): The current graph state

    Returns:
        str: Binary decision for next node to call
    """

    print("---ASSESS GRADED DOCUMENTS---")
    state["question"]
    filtered_documents = state["documents"]

    if not filtered_documents:
        # All documents have been filtered check_relevance
        # We will re-generate a new query
        print(
            "---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, TRANSFORM QUERY---"
        )
        return "transform_query"
    else:
        # We have relevant documents, so generate answer
        print("---DECISION: GENERATE---")
        return "generate"


def grade_generation_v_documents_and_question(state):
    """
    Determines whether the generation is grounded in the document and answers question.

    Args:
        state (dict): The current graph state

    Returns:
        str: Decision for next node to call
    """

    print("---CHECK HALLUCINATIONS---")
    question = state["question"]
    documents = state["documents"]
    generation = state["generation"]

    score = hallucination_grader.invoke(
        {"documents": documents, "generation": generation}
    )
    grade = score.binary_score

    # Check hallucination
    if grade == "yes":
        print("---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---")
        # Check question-answering
        print("---GRADE GENERATION vs QUESTION---")
        score = answer_grader.invoke({"question": question, "generation": generation})
        grade = score.binary_score
        if grade == "yes":
            print("---DECISION: GENERATION ADDRESSES QUESTION---")
            return "useful"
        else:
            print("---DECISION: GENERATION DOES NOT ADDRESS QUESTION---")
            return "not useful"
    else:
        print("---DECISION: GENERATION IS NOT GROUNDED IN DOCUMENTS, RE-TRY---")
        return "not supported"



from langgraph.graph import END, StateGraph, START
from patagent.agent.patent_rag import *
from typing_extensions import TypedDict
from typing import List


class GraphState(TypedDict):
    """
    Represents the state of our graph.

    Attributes:
        question: question
        generation: LLM generation
        documents: list of documents
    """
    question: str
    generation: str
    documents: List[str]

workflow = StateGraph(GraphState)

# Define the nodes
workflow.add_node("chitchat", chitchat)  # retriever
workflow.add_node("retrieve", retrieve)  # retriever
workflow.add_node("hybrid_retrieve", hybrid_retrieve)  # hybrid_retriever
workflow.add_node("grade_documents", grade_documents)  # grade documents
workflow.add_node("generate", generate)  # generatae
workflow.add_node("transform_query", transform_query)  # transform_query
workflow.add_node("request_data", request_data)

# Build graph
workflow.add_conditional_edges(
    START,
    route_question,
    {
        "chitchat": "chitchat",
        "retrieve": "retrieve",
        "hybrid_retrieve": "hybrid_retrieve",
        "data": "request_data"
    },
)
workflow.add_edge("retrieve", "grade_documents")
workflow.add_edge("hybrid_retrieve", "grade_documents")
workflow.add_conditional_edges(
    "grade_documents",
    decide_to_generate,
    {
        "transform_query": "transform_query",
        "generate": "generate",
    },
)
workflow.add_edge("transform_query", "retrieve")
workflow.add_edge("generate", END)
workflow.add_edge("request_data", END)
workflow.add_edge("chitchat", END)

# Compile
graph = workflow.compile()

# print(graph.invoke({'question': '你好'}))